from langchain.embeddings import CacheBackedEmbeddings
from langchain.storage import LocalFileStore
from langchain_community.embeddings import DashScopeEmbeddings
from langchain_community.vectorstores import FAISS


class FaissDbUtils:
    def __init__(self):
        self.embeddings = DashScopeEmbeddings()
        self.store = LocalFileStore("./cache/")
        self.cached_emb = CacheBackedEmbeddings.from_bytes_store(self.embeddings, self.store,
                                                            namespace=self.embeddings.model)

    def add(self, chunks, keyindex):
        db = FAISS.from_documents(chunks, self.cached_emb)
        db.save_local(keyindex)

    def search(self, keyindex, question, count):
        db = FAISS.load_local(keyindex, self.cached_emb,
                              allow_dangerous_deserialization=True)
        res = db.similarity_search(question, k=count)
        return res
